##这里数据采用了模拟的数据，模型使用了MLP，在做课程设计时应用真实数据代替，并采用合适的模型和评价指标
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import matplotlib.pyplot as plt
import os
from datetime import datetime
# 定义生成器模型
class ECGGenerator(nn.Module):
    def __init__(self, latent_dim=100, output_dim=256):
        super(ECGGenerator, self).__init__()
        self.latent_dim = latent_dim
        
        self.model = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.LeakyReLU(0.2),
            nn.Linear(128, 256),
            nn.LeakyReLU(0.2),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.Linear(512, output_dim),
            nn.Tanh()  # 输出在[-1,1]范围内
        )
    
    def forward(self, z):
        return self.model(z)

# 定义判别器模型
class ECGDiscriminator(nn.Module):
    def __init__(self, input_dim=256):
        super(ECGDiscriminator, self).__init__()
        
        self.model = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, 128),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )
    
    def forward(self, ecg):
        return self.model(ecg)

# 自定义ECG数据集
class ECGDataset(Dataset):
    def __init__(self, num_samples=1000, ecg_length=256):
        self.num_samples = num_samples
        self.ecg_length = ecg_length
        # 这里应该加载真实ECG数据，这里用模拟数据代替
        self.data = self._generate_synthetic_ecgs()
    
    def _generate_synthetic_ecgs(self):
        """生成简单的模拟ECG数据"""
        data = []
        for _ in range(self.num_samples):
            # 基础ECG波形
            x = np.linspace(0, 2*np.pi, self.ecg_length)
            ecg = np.sin(x) * 1.5
            
            # 添加P波
            p_start = int(0.2 * self.ecg_length)
            p_end = int(0.3 * self.ecg_length)
            ecg[p_start:p_end] += np.sin(np.linspace(0, np.pi, p_end-p_start)) * 0.5
            
            # 添加QRS复合波
            qrs_start = int(0.4 * self.ecg_length)
            qrs_end = int(0.55 * self.ecg_length)
            ecg[qrs_start:qrs_end] += np.sin(np.linspace(0, 3*np.pi, qrs_end-qrs_start)) * 2.0
            
            # 添加T波
            t_start = int(0.6 * self.ecg_length)
            t_end = int(0.8 * self.ecg_length)
            ecg[t_start:t_end] += np.sin(np.linspace(0, np.pi, t_end-t_start)) * 0.8
            
            # 添加噪声
            noise = np.random.normal(0, 0.1, self.ecg_length)
            ecg += noise
            
            # 标准化
            ecg = (ecg - ecg.min()) / (ecg.max() - ecg.min()) * 2 - 1  # 缩放到[-1,1]
            
            data.append(ecg)
        
        return torch.FloatTensor(np.array(data))
    
    def __len__(self):
        return self.num_samples
    
    def __getitem__(self, idx):
        return self.data[idx]

# 训练函数
def train_gan(generator, discriminator, dataloader, epochs=1000, latent_dim=100, device='cuda:0'):
    # 检查设备可用性
    if not torch.cuda.is_available() or not device.startswith('cuda'):
        device = 'cpu'
        print("CUDA not available, using CPU instead")
    
    device = torch.device(device)
    print(f"Using device: {device}")
    
    generator.to(device)
    discriminator.to(device)
    
    # 损失函数和优化器
    criterion = nn.BCELoss()
    lr = 0.0002
    g_optimizer = optim.Adam(generator.parameters(), lr=lr)
    d_optimizer = optim.Adam(discriminator.parameters(), lr=lr)
    
    for epoch in range(epochs):
        for i, real_ecgs in enumerate(dataloader):
            batch_size = real_ecgs.size(0)
            real_ecgs = real_ecgs.to(device)
            
            # 真实和假标签
            real_labels = torch.ones(batch_size, 1).to(device)
            fake_labels = torch.zeros(batch_size, 1).to(device)
            
            # 训练判别器
            d_optimizer.zero_grad()
            
            # 真实ECG的损失
            real_outputs = discriminator(real_ecgs)
            d_loss_real = criterion(real_outputs, real_labels)
            
            # 假ECG的损失
            z = torch.randn(batch_size, latent_dim).to(device)
            fake_ecgs = generator(z)
            fake_outputs = discriminator(fake_ecgs.detach())
            d_loss_fake = criterion(fake_outputs, fake_labels)
            
            # 总判别器损失
            d_loss = d_loss_real + d_loss_fake
            d_loss.backward()
            d_optimizer.step()
            
            # 训练生成器
            g_optimizer.zero_grad()
            
            z = torch.randn(batch_size, latent_dim).to(device)
            fake_ecgs = generator(z)
            outputs = discriminator(fake_ecgs)
            g_loss = criterion(outputs, real_labels)
            
            g_loss.backward()
            g_optimizer.step()
            
        if epoch % 1 == 0:
            print(f"Epoch [{epoch}/{epochs}], d_loss: {d_loss.item():.4f}, g_loss: {g_loss.item():.4f}")
            
            # 可视化生成的ECG
            with torch.no_grad():
                test_z = torch.randn(1, latent_dim).to(device)
                generated_ecg = generator(test_z).cpu().numpy()[0]
                
                plt.figure(figsize=(10, 4))
                plt.plot(generated_ecg)
                plt.title(f"Generated ECG at Epoch {epoch}")
                plt.show()
def save_ecg_plot(ecg_data, filename, title="Generated ECG"):
    """保存ECG图像到文件"""
    plt.figure(figsize=(10, 4))
    plt.plot(ecg_data)
    plt.title(title)
    plt.savefig(filename)
    plt.close()
# 主程序
if __name__ == "__main__":
    # 参数设置
    latent_dim = 100
    ecg_length = 256
    batch_size = 32
    epochs = 10
    device = 'cuda:0'  # 默认使用cuda:0
    
    # 创建模型
    generator = ECGGenerator(latent_dim, ecg_length)
    discriminator = ECGDiscriminator(ecg_length)
    
    # 创建数据集和数据加载器
    dataset = ECGDataset(num_samples=1000, ecg_length=ecg_length)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    # 训练模型
    train_gan(generator, discriminator, dataloader, epochs, latent_dim, device)
    
    # 训练后生成示例ECG
    with torch.no_grad():
        test_z = torch.randn(5, latent_dim).to(torch.device(device))
        generated_ecgs = generator(test_z).cpu().numpy()
        
        plt.figure(figsize=(15, 8))
        for i in range(5):
            plt.subplot(5, 1, i+1)
            plt.plot(generated_ecgs[i])
            plt.title(f"Generated ECG Sample {i+1}")
        plt.tight_layout()
        plt.show()
    with torch.no_grad():
        test_z = torch.randn(5, latent_dim).to(torch.device(device))
        generated_ecgs = generator(test_z).cpu().numpy()
        output_dir = './'  # 输出目录
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        for i in range(5):
            filename = os.path.join(output_dir, f"final_ecg_sample_{i+1}_{timestamp}.png")
            save_ecg_plot(generated_ecgs[i], filename, f"Final Generated ECG Sample {i+1}")
        
        # 保存一个包含所有样本的图
        plt.figure(figsize=(15, 8))
        for i in range(5):
            plt.subplot(5, 1, i+1)
            plt.plot(generated_ecgs[i])
            plt.title(f"Generated ECG Sample {i+1}")
        plt.tight_layout()
        plt.savefig(os.path.join(output_dir, f"all_samples_{timestamp}.png"))
        plt.close()